/**
 * @file eigen_common.h
 *
 * Common facilities for using Eigen library
 * 
 * @author Dahua Lin
 */

#ifndef SLIPP_EIGEN_COMMON_H_
#define SLIPP_EIGEN_COMMON_H_

#include <slipp/base/basic_defs.h>
#include <Eigen/Dense>
#include <algorithm>
#include <vector>

#define EIGEN_AROW(T) Eigen::Array<T, 1, Eigen::Dynamic>
#define EIGEN_ACOL(T) Eigen::Array<T, Eigen::Dynamic, 1>
#define EIGEN_ARR(T) Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>

#define EIGEN_VROW(T) Eigen::Matrix<T, 1, Eigen::Dynamic>
#define EIGEN_VCOL(T) Eigen::Matrix<T, Eigen::Dynamic, 1>
#define EIGEN_MAT(T) Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>

namespace sli
{
	/****************************************************************
	 *
	 *   Basics
	 *
	 ****************************************************************/

	// eval type helper

	template<class Derived>
	struct eigen_evaluated_type
	{
		typedef typename Eigen::internal::eval<Derived>::type type;
	};

	template<typename ScalarType, int Rows, int Cols>
	struct eigen_array
	{
		typedef Eigen::Array<ScalarType, Rows, Cols> type;
	};

	template<typename ScalarType, int Rows, int Cols>
	struct eigen_matrix
	{
		typedef Eigen::Matrix<ScalarType, Rows, Cols> type;
	};

	/****************************************************************
	 *
	 *   Mapping and Filtering
	 *
	 ****************************************************************/

	template<class Func, class Derived>
	inline Eigen::CwiseUnaryOp<Func, const Derived>
	map(Func func, const Eigen::DenseBase<Derived>& x)
	{
		return Eigen::CwiseUnaryOp<Func, const Derived>(x.derived(), func);
	}

	template<class Func, class Derived1, class Derived2>
	inline Eigen::CwiseBinaryOp<Func, const Derived1, const Derived2>
	map(Func func, const Eigen::DenseBase<Derived1>& x1, const Eigen::DenseBase<Derived2>& x2)
	{
		return Eigen::CwiseBinaryOp<Func, const Derived1, const Derived2>(x1.derived(), x2.derived(), func);
	}

	typedef Eigen::Array<int32_t, Eigen::Dynamic, 1> eigen_index_arr_t;


	template<class Derived>
	inline eigen_index_arr_t find(const Eigen::DenseBase<Derived>& x)
	{
		std::vector<int32_t> v;
		for (int32_t i = 0; i < x.size(); ++i)
		{
			if (x.coeff(i)) v.push_back(i);
		}

		eigen_index_arr_t r((index_t)v.size());
		std::copy_n(v.data(), v.size(), r.data());
		return r;
	}


	template<class Predicate, class Derived>
	inline eigen_index_arr_t find(Predicate func, const Eigen::DenseBase<Derived>& x)
	{
		std::vector<int32_t> v;
		for (int32_t i = 0; i < x.size(); ++i)
		{
			if (func(x.coeff(i))) v.push_back(i);
		}

		eigen_index_arr_t r((index_t)v.size());
		std::copy_n(v.data(), v.size(), r.data());
		return r;
	}

	template<class Predicate, class Derived>
	inline eigen_index_arr_t find(Predicate func, const Eigen::DenseBase<Derived>& x, const Eigen::DenseBase<Derived>& y)
	{
		std::vector<int32_t> v;
		for (int32_t i = 0; i < x.size(); ++i)
		{
			if (func(x.coeff(i), y.coeff(i))) v.push_back(i);
		}

		eigen_index_arr_t r((index_t)v.size());
		std::copy_n(v.data(), v.size(), r.data());
		return r;
	}


	/****************************************************************
	 *
	 *   Sub-array selection
	 *
	 ****************************************************************/

	// reader/writer classes

	namespace _detail
	{
		template<
			class Derived,
			class IdxDerived,
			template<typename ScalarType, int Rows, int Cols> class ResultT
		>
		class eigen_coeffs_reader
		{
		public:
			typedef typename Derived::Scalar scalar_type;
			typedef typename ResultT<scalar_type,
					IdxDerived::RowsAtCompileTime,
					IdxDerived::ColsAtCompileTime>::type result_type;

		public:
			eigen_coeffs_reader(const Derived& src, const IdxDerived& idx)
			: m_src(const_cast<Derived&>(src)), m_idx(idx)
			{
			}

			operator result_type() const
			{
				const typename eigen_evaluated_type<Derived>::type& es = m_src.eval();
				const typename eigen_evaluated_type<IdxDerived>::type& ei = m_idx.eval();

				result_type r(ei.rows(), ei.cols());

				index_t N = ei.size();
				for (index_t i = 0; i < N; ++i)
				{
					r.coeffRef(i) = es(ei.coeff(i));
				}

				return r;
			}

		protected:
			Derived& m_src;
			const IdxDerived& m_idx;
		};


		template<
			class Derived,
			class IdxDerived,
			template<typename ScalarType, int Rows, int Cols> class ResultT
		>
		class eigen_coeffs2_reader
		{
		public:
			typedef typename Derived::Scalar scalar_type;
			typedef typename ResultT<scalar_type,
					IdxDerived::RowsAtCompileTime,
					IdxDerived::ColsAtCompileTime>::type result_type;

		public:
			eigen_coeffs2_reader(const Derived& src, const IdxDerived& irow, const IdxDerived& icol)
			: m_src(const_cast<Derived&>(src)), m_irow(irow), m_icol(icol)
			{
			}

			operator result_type() const
			{
				const typename eigen_evaluated_type<Derived>::type& es = m_src.eval();
				const typename eigen_evaluated_type<IdxDerived>::type& eirow = m_irow.eval();
				const typename eigen_evaluated_type<IdxDerived>::type& eicol = m_icol.eval();

				result_type r(eirow.rows(), eicol.cols());

				index_t N = eirow.size();
				for (index_t i = 0; i < N; ++i)
				{
					r.coeffRef(i) = es(eirow.coeff(i), eicol.coeff(i));
				}

				return r;
			}

		protected:
			Derived& m_src;
			const IdxDerived& m_irow;
			const IdxDerived& m_icol;
		};


		template<
			class Derived,
			class IdxDerived,
			template<typename ScalarType, int Rows, int Cols> class ResultT
		>
		class eigen_rows_reader
		{
		public:
			typedef typename Derived::Scalar scalar_type;
			typedef typename ResultT<scalar_type,
					IdxDerived::SizeAtCompileTime,
					Derived::ColsAtCompileTime>::type result_type;

		public:
			eigen_rows_reader(const Derived& src, const IdxDerived& irow)
			: m_src(const_cast<Derived&>(src)), m_irow(irow)
			{
			}

			operator result_type() const
			{
				const typename eigen_evaluated_type<Derived>::type& es = m_src.eval();
				const typename eigen_evaluated_type<IdxDerived>::type& eirow = m_irow.eval();

				index_t m = eirow.size();
				index_t n = es.cols();
				result_type r(m, n);

				for (index_t j = 0; j < n; ++j)
				{
					for (index_t i = 0; i < m; ++i)
					{
						r.coeffRef(i, j) = es(eirow.coeff(i), j);
					}
				}

				return r;
			}

		protected:
			Derived& m_src;
			const IdxDerived& m_irow;
		};


		template<
			class Derived,
			class IdxDerived,
			template<typename ScalarType, int Rows, int Cols> class ResultT
		>
		class eigen_cols_reader
		{
		public:
			typedef typename Derived::Scalar scalar_type;
			typedef typename ResultT<scalar_type,
					Derived::RowsAtCompileTime,
					IdxDerived::SizeAtCompileTime>::type result_type;

		public:
			eigen_cols_reader(const Derived& src, const IdxDerived& icol)
			: m_src(const_cast<Derived&>(src)), m_icol(icol)
			{
			}

			operator result_type() const
			{
				const typename eigen_evaluated_type<Derived>::type& es = m_src.eval();
				const typename eigen_evaluated_type<IdxDerived>::type& eicol = m_icol.eval();

				index_t m = es.rows();
				index_t n = eicol.size();
				result_type r(m, n);

				for (index_t j = 0; j < n; ++j)
				{
					for (index_t i = 0; i < m; ++i)
					{
						r.coeffRef(i, j) = es(i, eicol.coeff(j));
					}
				}

				return r;
			}

		protected:
			Derived& m_src;
			const IdxDerived& m_icol;
		};


		template<
			class Derived,
			class IRowDerived,
			class IColDerived,
			template<typename ScalarType, int Rows, int Cols> class ResultT
		>
		class eigen_rows_and_cols_reader
		{
		public:
			typedef typename Derived::Scalar scalar_type;
			typedef typename ResultT<scalar_type,
					IRowDerived::SizeAtCompileTime,
					IColDerived::SizeAtCompileTime>::type result_type;

		public:
			eigen_rows_and_cols_reader(const Derived& src, const IRowDerived& irow, const IColDerived& icol)
			: m_src(const_cast<Derived&>(src)), m_irow(irow), m_icol(icol)
			{
			}

			operator result_type() const
			{
				const typename eigen_evaluated_type<Derived>::type& es = m_src.eval();
				const typename eigen_evaluated_type<IRowDerived>::type& eirow = m_irow.eval();
				const typename eigen_evaluated_type<IColDerived>::type& eicol = m_icol.eval();

				index_t m = eirow.size();
				index_t n = eicol.size();
				result_type r(m, n);

				for (index_t j = 0; j < n; ++j)
				{
					for (index_t i = 0; i < m; ++i)
					{
						r.coeffRef(i, j) = es(eirow.coeff(i), eicol.coeff(j));
					}
				}

				return r;
			}

		protected:
			Derived& m_src;
			const IRowDerived& m_irow;
			const IColDerived& m_icol;
		};

	}


	template<class Derived, class IdxDerived>
	inline _detail::eigen_coeffs_reader<Derived, IdxDerived, eigen_matrix>
	select_coeffs(const Eigen::MatrixBase<Derived>& x, const Eigen::DenseBase<IdxDerived>& idx)
	{
		return _detail::eigen_coeffs_reader<Derived, IdxDerived, eigen_matrix>(x.derived(), idx.derived());
	}

	template<class Derived, class IdxDerived>
	inline _detail::eigen_coeffs_reader<Derived, IdxDerived, eigen_array>
	select_coeffs(const Eigen::ArrayBase<Derived>& x, const Eigen::DenseBase<IdxDerived>& idx)
	{
		return _detail::eigen_coeffs_reader<Derived, IdxDerived, eigen_array>(x.derived(), idx.derived());
	}

	template<class Derived, class IdxDerived>
	inline _detail::eigen_coeffs2_reader<Derived, IdxDerived, eigen_matrix>
	select_coeffs(const Eigen::MatrixBase<Derived>& x,
			const Eigen::DenseBase<IdxDerived>& irow, const Eigen::DenseBase<IdxDerived>& icol)
	{
		return _detail::eigen_coeffs2_reader<Derived, IdxDerived, eigen_matrix>(x.derived(), irow.derived(), icol.derived());
	}

	template<class Derived, class IdxDerived>
	inline _detail::eigen_coeffs2_reader<Derived, IdxDerived, eigen_array>
	select_coeffs(const Eigen::ArrayBase<Derived>& x,
			const Eigen::DenseBase<IdxDerived>& irow, const Eigen::DenseBase<IdxDerived>& icol)
	{
		return _detail::eigen_coeffs2_reader<Derived, IdxDerived, eigen_array>(x.derived(), irow.derived(), icol.derived());
	}

	template<class Derived, class IdxDerived>
	inline _detail::eigen_rows_reader<Derived, IdxDerived, eigen_matrix>
	select_rows(const Eigen::MatrixBase<Derived>& x, const Eigen::DenseBase<IdxDerived>& idx)
	{
		return _detail::eigen_rows_reader<Derived, IdxDerived, eigen_matrix>(x.derived(), idx.derived());
	}

	template<class Derived, class IdxDerived>
	inline _detail::eigen_rows_reader<Derived, IdxDerived, eigen_array>
	select_rows(const Eigen::ArrayBase<Derived>& x, const Eigen::DenseBase<IdxDerived>& idx)
	{
		return _detail::eigen_rows_reader<Derived, IdxDerived, eigen_array>(x.derived(), idx.derived());
	}

	template<class Derived, class IdxDerived>
	inline _detail::eigen_cols_reader<Derived, IdxDerived, eigen_matrix>
	select_cols(const Eigen::MatrixBase<Derived>& x, const Eigen::DenseBase<IdxDerived>& idx)
	{
		return _detail::eigen_cols_reader<Derived, IdxDerived, eigen_matrix>(x.derived(), idx.derived());
	}

	template<class Derived, class IdxDerived>
	inline _detail::eigen_cols_reader<Derived, IdxDerived, eigen_array>
	select_cols(const Eigen::ArrayBase<Derived>& x, const Eigen::DenseBase<IdxDerived>& idx)
	{
		return _detail::eigen_cols_reader<Derived, IdxDerived, eigen_array>(x.derived(), idx.derived());
	}

	template<class Derived, class IRowDerived, class IColDerived>
	inline _detail::eigen_rows_and_cols_reader<Derived, IRowDerived, IColDerived, eigen_matrix>
	select_rows_cols(const Eigen::MatrixBase<Derived>& x,
			const Eigen::DenseBase<IRowDerived>& irow, const Eigen::DenseBase<IColDerived>& icol)
	{
		return _detail::eigen_rows_and_cols_reader<Derived, IRowDerived, IColDerived, eigen_matrix>(
				x.derived(), irow.derived(), icol.derived());
	}

	template<class Derived, class IRowDerived, class IColDerived>
	inline _detail::eigen_rows_and_cols_reader<Derived, IRowDerived, IColDerived, eigen_array>
	select_rows_cols(const Eigen::ArrayBase<Derived>& x,
			const Eigen::DenseBase<IRowDerived>& irow, const Eigen::DenseBase<IColDerived>& icol)
	{
		return _detail::eigen_rows_and_cols_reader<Derived, IRowDerived, IColDerived, eigen_array>(
				x.derived(), irow.derived(), icol.derived());
	}

}


#endif 
